import os
import shutil
import random
from os import path as osp

random.seed(2024)

def copy_files(ori_raw_path, ori_annos_path, dst_raw_path, dst_annos_path, num_val):
    # 确保验证集和训练集的文件夹存在
    val_raw_path = osp.join(dst_raw_path, "val", "normal")
    val_annos_path = osp.join(dst_annos_path, "val", "normal")
    train_raw_path = osp.join(dst_raw_path, "train", "normal")
    train_annos_path = osp.join(dst_annos_path, "train", "normal")

    os.makedirs(val_raw_path, exist_ok=True)
    os.makedirs(val_annos_path, exist_ok=True)
    os.makedirs(train_raw_path, exist_ok=True)
    os.makedirs(train_annos_path, exist_ok=True)
    raw_files = []
    for root, _, files in os.walk(ori_raw_path):
        for file in files:
            if file.endswith(".raw"):
                raw_files.append(osp.join(root, file))

    random.shuffle(raw_files)
    val_count = 0

    for i, raw_file in enumerate(raw_files):
        file_name = osp.basename(raw_file)
        base_name = osp.splitext(file_name)[0]
        cor_annos_path = osp.join(ori_annos_path, base_name, "annotations")

        if not osp.exists(cor_annos_path):
            print(f"Annotations not found for {base_name}")
            continue

        if val_count < num_val:
            dst_raw_subpath = osp.join(dst_raw_path, "val", "normal")
            dst_annos_subpath = osp.join(dst_annos_path, "val", "normal", base_name,"annotation")
        else:
            dst_raw_subpath = osp.join(dst_raw_path, "train", "normal")
            dst_annos_subpath = osp.join(dst_annos_path, "train", "normal", base_name,"annotation")

        if not osp.exists(dst_raw_subpath):
            os.makedirs(dst_raw_subpath)
        if not osp.exists(dst_annos_subpath):
            os.makedirs(dst_annos_subpath)

        # 复制 raw 文件
        shutil.copy(raw_file, dst_raw_subpath)
            
        # 复制 annotations 文件夹中的 txt 文件
        for ann_file in os.listdir(cor_annos_path):
            if ann_file.endswith(".txt"):
                full_ann_file = osp.join(cor_annos_path, ann_file)
                if osp.isfile(full_ann_file):
                    shutil.copy(full_ann_file, dst_annos_subpath)

        if val_count < num_val:
            val_count += 1

        print(f"Copied {file_name} to {dst_raw_subpath} and its annotations to {dst_annos_subpath}")

if __name__ == "__main__":
    ori_raw_path = "E:\\Datasets\\eventraw"
    ori_annos_path = "E:\\Datasets\\eventframe"
    dst_raw_path = "E:\\Datasets\\lastdata\\raw"
    dst_annos_path = "E:\\Datasets\\lastdata\\annotations"
    num_val = 0  # 设置验证集的数量

    copy_files(ori_raw_path, ori_annos_path, dst_raw_path, dst_annos_path, num_val)